package com.we.yarn;

import org.apache.flink.util.concurrent.ExecutorThreadFactory;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.yarn.api.records.ApplicationId;
import org.apache.hadoop.yarn.api.records.ApplicationReport;
import org.apache.hadoop.yarn.api.records.YarnApplicationState;
import org.apache.hadoop.yarn.client.api.YarnClient;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.apache.hadoop.yarn.exceptions.YarnException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.text.SimpleDateFormat;
import java.util.*;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;

/**
 * @author JasonLee
 * @description 获取yarn队列上面任务的状态信息
 * @date $date$ $time$
 */
public class YarnClientUtils {

    private static final Logger LOG = LoggerFactory.getLogger(YarnClientUtils.class);
    private static YarnClient yarnClient;
    private static final SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
    private static transient ScheduledExecutorService scheduler =
            Executors.newScheduledThreadPool(1, new ExecutorThreadFactory("alarm-connector-pool"));;

    private static Map<String, ApplicationReport> applicationReportMap = new HashMap<>(200);

    static {
        Configuration conf = new YarnConfiguration();
        yarnClient = YarnClient.createYarnClient();
        yarnClient.init(conf);
        yarnClient.start();
    }

    public static Map<String, ApplicationReport> getApplicationReport()
            throws IOException, YarnException {
        Map<String, ApplicationReport> map = new HashMap<>();
        // 获取running状态任务的集合
        List<ApplicationReport> applications =
                yarnClient.getApplications(EnumSet.of(YarnApplicationState.RUNNING));
        for (ApplicationReport applicationReport : applications) {
            map.put(applicationReport.getName(), applicationReport);
        }
        LOG.info("获取 Yarn 任务状态成功 : {}", map);
        return map;
    }

    public static List<ApplicationReport> getApplicationReportList()
            throws IOException, YarnException {
        List<ApplicationReport> list = new ArrayList<>();
        // 获取running状态任务的集合
        List<ApplicationReport> applications =
                yarnClient.getApplications(EnumSet.of(YarnApplicationState.RUNNING));
        for (ApplicationReport applicationReport : applications) {
            list.add(applicationReport);
        }
        LOG.info("获取 Yarn 任务成功,正在运行的任务有 : {} 个", list.size());
        return list;
    }

    public static void main(String[] args) throws IOException, YarnException {
        final List<ApplicationReport> applicationReportList = getApplicationReportList();
        applicationReportList.forEach(
                applicationReport -> {
                    System.out.println(applicationReport.getApplicationId());
                    System.out.println(applicationReport.getName());
                    System.out.println(applicationReport.getStartTime());
                    System.out.println(applicationReport.getTrackingUrl());
                    System.out.println(applicationReport.getYarnApplicationState());
                });
    }

    /**
     * 判断任务是否挂了
     *
     * @param applicationId
     * @return
     */
    public static boolean isDown(ApplicationId applicationId) {
        try {
            ApplicationReport applicationReport = yarnClient.getApplicationReport(applicationId);
            YarnApplicationState yarnApplicationState = applicationReport.getYarnApplicationState();
            if (YarnApplicationState.FINISHED.equals(yarnApplicationState)) {
                return true;
            }
        } catch (YarnException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return false;
    }

    /**
     * 根据applicationid删除任务
     *
     * @param applicationId
     * @throws IOException
     * @throws YarnException
     */
    public static void deleteApplication(ApplicationId applicationId)
            throws IOException, YarnException {
        yarnClient.killApplication(applicationId);
    }
}
